package cloud.keveri.example.chat4j.interfaces;

import cloud.keveri.chat4j.api.ChatBlender;
import cloud.keveri.chat4j.api.Response;
import cloud.keveri.chat4j.api.entity.chat.CompletionRequest;
import cloud.keveri.chat4j.api.entity.chat.Message;
import cloud.keveri.chat4j.api.entity.image.ImageRequest;
import cloud.keveri.chat4j.core.factory.ChatFactory;
import cloud.keveri.chat4j.provider.SupplierType;
import cloud.keveri.chat4j.provider.chatgpt.listener.DefaultChatgptEventSourceListener;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import javax.servlet.http.HttpServletResponse;
import java.util.ArrayList;
import java.util.List;

/**
 *  测试处理器
 * <p>
 * 作者：沈自在
 **/
@RestController
public class TestController {

    @PostMapping("/api/chatgpt/completion/not")
    public Response chatgpt_completion_not_stream() {

        ChatBlender chatBlender = ChatFactory.acquireChatBlender(SupplierType.CHATGPT);

        Message message = new Message("user", "讲个笑话吧");
        List<Message> messages = new ArrayList<>();
        messages.add(message);

        CompletionRequest request = CompletionRequest.builder()
                .messages(messages)
                .build();

        Response response = chatBlender.completions(request);
        System.out.println("response = " + response);
        return response;
    }

    @GetMapping("/api/chatgpt/completion/sse")
    public ResponseBodyEmitter sseEmitter(String question, HttpServletResponse response) {
        // 1. 准备请求

        response.setContentType("text/event-stream");
        response.setCharacterEncoding("UTF-8");
        response.setHeader("Cache-Control", "no-cache");


        Message message = Message.of(question);
        List<Message> messages = new ArrayList<>();
        messages.add(message);

        CompletionRequest request = CompletionRequest.builder()
                .messages(messages)
                .stream(true)
                .build();

        // 2. sse
        ResponseBodyEmitter emitter = new ResponseBodyEmitter(3 * 60 * 1000L);
        DefaultChatgptEventSourceListener listener = new DefaultChatgptEventSourceListener(emitter);
        //ConsoleChatgptEventSourceListener listener = new ConsoleChatgptEventSourceListener();

        // 3. sdk
        ChatBlender chatBlender = ChatFactory.acquireChatBlender(SupplierType.CHATGPT);
        chatBlender.completions(request, listener);

        return emitter;
    }

    @GetMapping("/api/chatgpt/image")
    public Response image(String prompt) {
        ChatBlender chatBlender = ChatFactory.acquireChatBlender(SupplierType.CHATGPT);
        ImageRequest request = ImageRequest.builder()
                .prompt(prompt)
                .build();
        return chatBlender.generateImage(request);
    }
}